package com.example.demo.arithmetic.datastructure.tree;

import lombok.Data;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * 哈夫曼树
 * 1.构建树的过程
 * （1）统计字符出现的频率，放入优先级队列（频率越高的编码后字符越短，这样可以节省空间）
 * （2）每次出队两个频次最低的元素，给他找个爹
 * （3）把爹重新放入队列，重复步骤2 3
 * （4）当队列中只有一个元素时候，树构建完成
 */
public class HuffmanTree {

    Node root;

    String str;

    Map<Character, Node> map = new HashMap<>();

    @Data
    static class Node {
        Character ch;
        int freq;
        Node left;
        Node right;
        String code;

        public Node(Character ch) {
            this.ch = ch;
        }

        public Node(int freq, Node left, Node right) {
            this.freq = freq;
            this.left = left;
            this.right = right;
        }

        boolean isLeaf() {
            return left == null;
        }
    }

    public HuffmanTree(String str) {
        this.str = str;
        char[] array = str.toCharArray();
        for (char c : array) {
            if (!map.containsKey(c)) {
                map.put(c, new Node(c));
            }
            Node node = map.get(c);
            node.freq++;

        }

        //        for (Node node : map.values()) {
        //            System.out.println(node);
        //        }
        PriorityQueue<Node> priorityQueue = new PriorityQueue<>(Comparator.comparingInt(Node::getFreq));
        priorityQueue.addAll(map.values());
        //
        // System.out.println(priorityQueue);
        while (priorityQueue.size() >= 2) {
            Node x = priorityQueue.poll();
            Node y = priorityQueue.poll();
            int f = x.freq + y.freq;

            Node dad = new Node(f, x, y);
            priorityQueue.offer(dad);
        }
        Node root = priorityQueue.poll();
        this.root = root;
        //        System.out.println(root);
        // 深度优先遍历
        //
        int sum = dfs(root, new StringBuilder());

//        System.out.println(sum);
//        for (Node node : map.values()) {
//            System.out.println(node);
//        }
    }


    public String encode() {
        StringBuilder sb = new StringBuilder();
        char[] chars = str.toCharArray();
        for (char ch : chars) {
            Node node = map.get(ch);
            sb.append(node.getCode());
        }
        return sb.toString();
    }

    public String decode(String decoded) {
        // 从根节点 寻找对应的字符
        //       数字是0向左走
        // 数字是1向右走
        // 如果没走到头 每走一步数字的索引 i++
        // 走到头就可以找到解码字符 再将node重置为根节点
        //
        char[] array = decoded.toCharArray();
        int i = 0;
        StringBuilder sb = new StringBuilder();
        Node node = root;
        while (i < array.length) {
            if (!node.isLeaf()) {
                if (array[i] == '0') {
                    node = node.left;
                } else if (array[i] == '1') {
                    node = node.right;
                }
                i++;
            }

            // 叶子节点
            if (node.isLeaf()) {
                sb.append(node.ch);
                node = root;
            }
        }
        return sb.toString();
    }

    /**
     * @param root
     * @param sb
     * @return 编码后的字节长度
     */
    private int dfs(Node root, StringBuilder sb) {
        int sum = 0;
        if (root.isLeaf()) {
            root.code = sb.toString();
            sum = root.freq * sb.length();
            //System.out.println(root + "" + sb);
            //            return;
        } else {
            sum += dfs(root.left, sb.append("0"));
            sum += dfs(root.right, sb.append("1"));
        }
        if (sb.length() > 0) {
            sb.deleteCharAt(sb.length() - 1);
        }
        return sum;
    }

    public static void main(String[] args) {
        HuffmanTree tree = new HuffmanTree("abbccccccc");
        String encode = tree.encode();
        System.out.println(encode);
        System.out.println(tree.decode(encode));
    }
}
